#include "hplanner_ompl/sPRM.h"
#include "ompl/base/GoalSampleableRegion.h"
#include "ompl/datastructures/NearestNeighborsSqrtApprox.h"
#include <algorithm>
#include <queue>
#include <limits>

namespace ompl
{
namespace geometric
{


static const unsigned int FIND_VALID_STATE_ATTEMPTS_WITHOUT_TIME_CHECK_SPRM = 10;

static const unsigned int BATCH_SIZE_SPRM = 200;


sPRM::sPRM(const base::SpaceInformationPtr& si) : base::Planner(si, "sPRM")
{
  specs_.recognizedGoal = base::GOAL_SAMPLEABLE_REGION;

  maxNearestNeighbors_ = 10;
  componentCount_ = 0;
  lastStart_ = NULL;
  lastGoal_ = NULL;
  nLastMilestones_ = 0;

  nPointCollisions_ = 0;
  nPathCollisions_ = 0;
  nCollisions_ = 0;
}

sPRM::~sPRM()
{
  freeMemory();
}

void sPRM::setProblemDefinition(const base::ProblemDefinitionPtr& pdef)
{
  Planner::setProblemDefinition(pdef);
  startM_.clear();
  goalM_.clear();
}

void sPRM::setMaxNearestNeighbors(unsigned int maxNearestNeighbors)
{
  maxNearestNeighbors_ = maxNearestNeighbors;
}

unsigned int sPRM::getMaxNearestNeighbors() const
{
  return maxNearestNeighbors_;
}

void sPRM::getPlannerData(base::PlannerData& data) const
{
  Planner::getPlannerData(data);

  for(unsigned int i = 0 ; i < milestones_.size() ; ++i)
    for(unsigned int j = 0 ; j < milestones_[i]->adjacent.size() ; ++j)
    {
      data.recordEdge(milestones_[i]->state, milestones_[i]->adjacent[j]->state);
      data.tagState(milestones_[i]->state, milestones_[i]->component);
    }
}

void sPRM::clear()
{
  /*
  Planner::clear();
  sampler_.reset();
  startM_.clear();
  goalM_.clear();
  lastStart_ = NULL;
  lastGoal_ = NULL;
  */

  Planner::clear();
  sampler_.reset();
  freeMemory();
  if (nn_)
      nn_->clear();
  milestones_.clear();
  startM_.clear();
  goalM_.clear();
  componentCount_ = 0;
  componentSizes_.clear();
  lastStart_ = NULL;
  lastGoal_ = NULL;
  nLastMilestones_ = 0;

  nPointCollisions_ = 0;
  nPathCollisions_ = 0;
  nCollisions_ = 0;
}

void sPRM::deepClear()
{
  Planner::clear();
  sampler_.reset();
  freeMemory();
  if (nn_)
      nn_->clear();
  milestones_.clear();
  startM_.clear();
  goalM_.clear();
  componentCount_ = 0;
  componentSizes_.clear();
  lastStart_ = NULL;
  lastGoal_ = NULL;
  nLastMilestones_ = 0;

  nPointCollisions_ = 0;
  nPathCollisions_ = 0;
  nCollisions_ = 0;
}

void sPRM::setup()
{
  Planner::setup();
  if (!nn_)
    nn_.reset(new NearestNeighborsSqrtApprox<Milestone*>());
  nn_->setDistanceFunction(boost::bind(&sPRM::distanceFunction, this, _1, _2));
}

bool sPRM::solve(const base::PlannerTerminationCondition& ptc)
{
  checkValidity();
  base::GoalSampleableRegion* goal = dynamic_cast<base::GoalSampleableRegion*>(pdef_->getGoal().get());

  if(!goal)
  {
    msg_.error("Goal undefined or unknown type of goal");
    return false;
  }

  unsigned int nrStartStates = milestones_.size();

  // add the valid start states as milestones
  while (const base::State *st = pis_.nextStart())
    startM_.push_back(addMilestone(si_->cloneState(st)));

  if(startM_.size() == 0)
  {
    msg_.error("There are no valid initial states!");
    return false;
  }

  if(!goal->canSample())
  {
    msg_.error("Insufficient states in sampleable goal region");
    return false;
  }

  if(!sampler_)
    sampler_ = si_->allocStateSampler();

  msg_.inform("Starting with %u states", nrStartStates);

  base::State* xstate = si_->allocState();
  std::pair<Milestone*, Milestone*> solEndpoints;

  while(ptc() == false)
  {
    if(goal->maxSampleCount() > goalM_.size() || goalM_.empty())
    {
      const base::State* st = goalM_.empty() ? pis_.nextGoal(ptc) : pis_.nextGoal();
      if(st)
        goalM_.push_back(addMilestone(si_->cloneState(st)));
      if(goalM_.empty())
      {
        msg_.error("Unable to find any valid goal states");
        break;
      }
    }

    if(haveSolution(startM_, goalM_, &solEndpoints))
    {
      constructSolution(solEndpoints.first, solEndpoints.second);
      break;
    }
    else
    {
      if(goal->maxSampleCount() > goalM_.size())
        growRoadmap(startM_, goalM_, base::PlannerOrTerminationCondition(ptc, base::timedPlannerTerminationCondition(0.1)), xstate);
      else
        growRoadmap(startM_, goalM_, ptc, xstate);


      if(haveSolution(startM_, goalM_, &solEndpoints))
      {
        constructSolution(solEndpoints.first, solEndpoints.second);
        break;
      }
    }
  }

  si_->freeState(xstate);
  msg_.inform("Created %u states", milestones_.size() - nrStartStates);

  msg_.inform("Perform %u point collisions", nPointCollisions_);
  msg_.inform("Perform %u path collisions", nPathCollisions_);
  msg_.inform("Perform %u collisions", nCollisions_);


  return goal->isAchieved();
}

void sPRM::reconstructLastSolution()
{
  if(lastStart_ && lastGoal_)
    constructSolution(lastStart_, lastGoal_);
  else
    msg_.warn("There is no solution to reconstruct");
}

void sPRM::freeMemory()
{
  for(unsigned int i = 0 ; i < milestones_.size() ; ++i)
  {
    if (milestones_[i]->state)
      si_->freeState(milestones_[i]->state);
    delete milestones_[i];
  }
}

bool sPRM::haveSolution(const std::vector<Milestone*> &start, const std::vector<Milestone*> &goal, std::pair<Milestone*, Milestone*> *endpoints)
{
  base::Goal* g = pdef_->getGoal().get();
  for(unsigned int i = 0; i < goal.size(); ++i)
  {
    for(unsigned int j = 0; j < start.size(); ++j)
    {
      if(goal[i]->component == start[j]->component && g->isStartGoalPairValid(goal[i]->state, start[j]->state))
      {
        if(endpoints)
        {
          endpoints->first = start[j];
          endpoints->second = goal[i];
        }
        return true;
      }
    }
  }

  return false;
}

void sPRM::constructSolution(const Milestone* start, const Milestone* goal)
{
  const unsigned int N = milestones_.size();
  std::vector<double> dist(N, std::numeric_limits<double>::infinity());
  std::vector<int> prev(N, -1);
  std::vector<int> seen(N, 0); // in fact extended

  dist[goal->index] = 0.0;
  for(unsigned int i = 0; i < N; ++i)
  {
    double minDist = std::numeric_limits<double>::infinity();
    int index = -1;
    for(unsigned int j = 0; j < N; ++j)
    {
      if(seen[j] == 0 && dist[j] < minDist)
      {
        minDist = dist[j];
        index = j;
      }
    }

    if(index < 0) break;
    seen[index] = 1;

    for(unsigned int j = 0; j < milestones_[index]->adjacent.size(); ++j)
    {
      const unsigned int idx = milestones_[index]->adjacent[j]->index;
      double altDist = dist[index] + milestones_[index]->costs[j];
      if(altDist < dist[idx])
      {
        dist[idx] = altDist;
        prev[idx] = index;
      }
    }
  }

  if(prev[start->index] >= 0)
  {
    PathGeometric* p = new PathGeometric(si_);
    int pos = start->index;
    do
    {
      p->states.push_back(si_->cloneState(milestones_[pos]->state));
      pos = prev[pos];
    }
    while(pos >= 0);
    pdef_->getGoal()->setSolutionPath(base::PathPtr(p));
    lastStart_ = start;
    lastGoal_ = goal;
  }
  else
      throw Exception(name_, "Internal error in computing shortest path");

}

double sPRM::distanceFunction(const Milestone* a, const Milestone* b) const
{
  return si_->distance(a->state, b->state);
}

bool sPRM::checkMotion(const base::State* s1, const base::State* s2) const
{
  nPathCollisions_++;

  {
    bool res = si_->isValid(s2);
    nCollisions_++;
    if(!res) return false;
  }

  base::StateSpace* stateSpace = si_->getStateSpace().get();

  bool result = true;
  int nd = stateSpace->validSegmentCount(s1, s2);

  std::queue<std::pair<int, int> > pos;
  if(nd >= 2)
  {
    pos.push(std::make_pair(1, nd - 1));

    base::State *test = si_->allocState();

    while (!pos.empty())
    {
      std::pair<int, int> x = pos.front();

      int mid = (x.first + x.second) / 2;
      stateSpace->interpolate(s1, s2, (double)mid / (double)nd, test);

      {
        bool res = si_->isValid(test);
        nCollisions_++;
        if(!res)
        {
          result = false;
          break;
        }
      }

      pos.pop();

      if(x.first < mid)
        pos.push(std::make_pair(x.first, mid - 1));
      if(x.second > mid)
        pos.push(std::make_pair(mid + 1, x.second));
    }

    si_->freeState(test);
  }

  return result;
}

sPRM::Milestone* sPRM::addMilestoneOnly(base::State *state)
{
  Milestone *m = new Milestone();
  m->state = state;
  m->component = componentCount_;
  componentSizes_[m->component] = 1;

  if(m->component == componentCount_)
    componentCount_++;

  m->index = milestones_.size();
  milestones_.push_back(m);
  nn_->add(m);
  return m;
}

sPRM::Milestone* sPRM::addMilestone(base::State *state)
{
  Milestone *m = new Milestone();
  m->state = state;
  m->component = componentCount_;
  componentSizes_[m->component] = 1;

  std::vector<Milestone*> nbh;
  nearestNeighbors(m, nbh);

  for(unsigned int i = 0; i < nbh.size(); ++i)
  {
    if(checkMotion(m->state, nbh[i]->state))
    {
      m->adjacent.push_back(nbh[i]);
      nbh[i]->adjacent.push_back(m);
      m->costs.push_back(si_->distance(m->state, nbh[i]->state));
      nbh[i]->costs.push_back(m->costs.back());
      uniteComponents(m, nbh[i]);
    }
  }

  if(m->component == componentCount_)
    componentCount_++;

  m->index = milestones_.size();
  milestones_.push_back(m);
  nn_->add(m);
  return m;
}

void sPRM::uniteComponents(Milestone *m1, Milestone *m2)
{
  if(m1->component == m2->component)
    return;

  if(m1->component < m2->component)
    std::swap(m1, m2);

  const unsigned long c = m2->component;
  componentSizes_[c] += componentSizes_[m1->component];
  componentSizes_.erase(m1->component);

  std::queue<Milestone*> q;
  q.push(m1);

  while(!q.empty())
  {
    Milestone* m = q.front();
    q.pop();
    if(m->component != c)
    {
      m->component = c;
      for(unsigned int i = 0; i < m->adjacent.size(); ++i)
        if(m->adjacent[i]->component != c)
          q.push(m->adjacent[i]);
    }
  }
}

void sPRM::nearestNeighbors(Milestone* milestone, std::vector<Milestone*>& nbh)
{
  nn_->nearestK(milestone, maxNearestNeighbors_, nbh);
}


void sPRM::updateConnectivity()
{
  for(unsigned int i = 0; i < milestones_.size(); ++i)
  {
    Milestone* m = milestones_[i];
    std::vector<Milestone*> nbh;
    nearestNeighbors(m, nbh);

    for(unsigned int j = 0; j < nbh.size(); ++j)
    {
      if(m->index < nLastMilestones_ && nbh[j]->index < nLastMilestones_)
        continue;

      if(m->index > nbh[j]->index)
      {
        if(checkMotion(m->state, nbh[j]->state))
        {
          m->adjacent.push_back(nbh[j]);
          nbh[j]->adjacent.push_back(m);
          m->costs.push_back(si_->distance(m->state, nbh[j]->state));
          nbh[j]->costs.push_back(m->costs.back());

          unsigned int old_m = m->component;
          unsigned int old_nbh = nbh[j]->component;
          uniteComponents(m, nbh[j]);

          if(m->component != old_m || old_nbh != nbh[j]->component)
          {
            componentCount_--;
          }
        }
      }
    }
  }
}


void sPRM::growRoadmap(const std::vector<Milestone*>& start,
                       const std::vector<Milestone*>& goal,
                       const base::PlannerTerminationCondition& ptc,
                       base::State* workState)
{
  nLastMilestones_ = milestones_.size();

  unsigned int new_milestone_num = 0;
  while(ptc() == false)
  {
    bool found = false;
    while(!found && ptc() == false)
    {
      unsigned int attempts = 0;
      do
      {
        sampler_->sampleUniform(workState);
        found = si_->isValid(workState);
        nPointCollisions_++;
        nCollisions_++;

        attempts++;
      }
      while(attempts < FIND_VALID_STATE_ATTEMPTS_WITHOUT_TIME_CHECK_SPRM && !found);
    }

    if(found)
    {
      addMilestoneOnly(si_->cloneState(workState));
      new_milestone_num++;
      if(new_milestone_num == BATCH_SIZE_SPRM)
        break;
    }
  }

  updateConnectivity();

}





}

}
